# 导入必要的模块
import asyncio
import base64
import datetime
import json
import os.path
import time

import cv2
from channels.generic.websocket import AsyncWebsocketConsumer
from cryptography.fernet import Fernet


from dataQuery.views import contraband_add
# 4
from yolov5.apply2 import Detect
import yaml
os.environ["DJANGO_ALLOW_ASYNC_UNSAFE"] = "true"

def get_yaml_data(yaml_file):
    with open(yaml_file, encoding='utf-8') as file:
        content = file.read()
        # 设置Loader=yaml.FullLoader忽略YAMLLoadWarning警告
        data = yaml.load(content, Loader=yaml.FullLoader)
    return data


config = get_yaml_data('./config/config.yaml')
if not os.path.exists(config['video_path']):
    os.makedirs(config['video_path'])

if not os.path.exists(config['image_path']):
    os.makedirs(config['image_path'])

# videoSet = VideoViewSet()
# detect = Detect()
# 定义 WebSocket 消费者类，继承自 AsyncWebsocketConsumer
# async def get_data():
#     return gen_display()
def image_to_base64(image):
    # 将图像数据编码为 Base64 字符串
    image_base64 = base64.b64encode(image).decode('utf-8')
    return image_base64


class TailfConsumer(AsyncWebsocketConsumer):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        print('init')
        self.detect = Detect()
        self.tasks = {}
        self.isRunning = True
        # self.camera = cv2.VideoCapture(0)
        self.camera = cv2.VideoCapture('D:\\data\\机场物流项目\\过机视频3.mp4')
        fps = 5.0
        # self.camera.set(cv2.CAP_PROP_FPS, fps)
        self.now_time = datetime.datetime.now()
        # 将当前时间转换为通用的日期和时间格式
        formatted_time = self.now_time.strftime("%Y_%m_%d %H_%M_%S")
        self.filename = '过机视频_' + str(formatted_time) + '.avi'
        #
        fourcc = cv2.VideoWriter_fourcc(*'XVID')
        self.out = cv2.VideoWriter(os.path.join(config['video_path'], self.filename), fourcc,20.0,
                                       (828,1024))

    async def connect(self):
        """
        这里定义连接建立时的逻辑
        """
        # 接受客户端连接
        await self.accept()

        # 开启循环发送数据
        self.isRunning = True
        print('begin')
        asyncio.ensure_future(self.send_data_loop())


    async def disconnect(self, close_code):
        # 连接断开时的逻辑

        self.isRunning = False
        self.out.release()
        # self.camera.release()
        # response = videoSet.createByrealtime(self.filename,os.path.join(config['video_path'], self.filename),self.now_time)
        # print(response)
        print('断开连接')

    async def send_data_loop(self):
        # camera = cv2.VideoCapture("D:/BaiduNetdiskDownload/普通过货视频20231106_095706/20231106_095706.mp4")


        while self.isRunning:
            # 获取数据

            data = self.gen_display()
            # frame_counter += 1
            # frameSize = camera.get(cv2.CAP_PROP_FRAME_COUNT)
            # if frame_counter == frameSize:
            #     frame_counter = 0
            #     camera.set(cv2.CAP_PROP_POS_FRAMES, 0)
            # 将数据转换为 JSON 字符串并发送给客户端

            await self.send(text_data=data)
            # 假设每秒发送一次数据
            await asyncio.sleep(0.2)

    def gen_display(self):
        """
        视频流生成器功能。
        """

        # 读取图片
        ret, frame = self.camera.read()

        if ret:

            result_frame, label_dicts = self.detect.detect(frame)
            self.out.write(result_frame)
            for label_dict in label_dicts:
                now_time = datetime.datetime.now()
                time_label = now_time.strftime("%Y_%m_%d %H_%M_%S")
                original_file = os.path.join(config['image_path'],
                                             str(label_dict['label']) + '_' + time_label + '_original.jpg')
                detect_file = os.path.join(config['image_path'],
                                           str(label_dict['label']) + '_' + time_label + '_detect.jpg')
                cv2.imwrite(original_file, frame)
                cv2.imwrite(detect_file, result_frame)
                unix_timestamp = int(time.mktime(now_time.timetuple()))
                contraband_add(unix_timestamp, config['place'], label_dict['label'], original_file, detect_file)
                label_dict['crop'] = image_to_base64(cv2.imencode('.jpeg', label_dict['crop'])[1])

            # frame = cv2.resize(frame, (0, 0), fx=0.25, fy=0.25)
            # 将图片进行解码
            ret, frame = cv2.imencode('.jpeg', result_frame)

            if ret:
                # 转换为byte类型的，存储在迭代器中
                # 遍历所有客户端连接，发送视频数据
                # response = HttpResponse(image_data, content_type='image/jpeg')
                encoded_data = json.dumps({"result_img": image_to_base64(frame), 'label_dicts': label_dicts})
                return encoded_data


class EncryptFile:
    # fernet_key = Fernet.generate_key()
    fernet_key = b'GyPq_tY3W9ldTG5Ln7cNrKI9EPf6GjQisssCk0qXyM0='
    cipher_suite = Fernet(fernet_key)

    def encrypt(self, file):
        with open(file, 'rb') as f:
            content = f.read()
            # 加密文件内容
            encrypted_content = self.cipher_suite.encrypt(content)
            # 创建加密文件的内容对象

        with open(file, 'wb') as f:
            f.write(encrypted_content)


encrypt = EncryptFile()

if __name__ == '__main__':
    encrypt.encrypt('./1.jpg')
